"""
计数模块
实现精确的进出计数逻辑
"""

import cv2
import numpy as np
from typing import Dict, List, Tuple, Optional, Set
from dataclasses import dataclass
from enum import Enum
import logging
from collections import defaultdict
import time

from tracker import TrackInfo
from region_manager import RegionManager

class PersonState(Enum):
    """人员状态枚举"""
    OUTSIDE = "outside"  # 在区域外
    INSIDE = "inside"    # 在区域内
    UNKNOWN = "unknown"  # 未知状态

@dataclass
class CountRecord:
    """计数记录"""
    track_id: int
    timestamp: float
    action: str  # "enter" 或 "exit"
    position: Tuple[int, int]
    confidence: float

class PersonCounter:
    """人员计数器"""
    
    def __init__(self, region_manager: RegionManager):
        """
        初始化计数器
        
        Args:
            region_manager: 区域管理器
        """
        self.region_manager = region_manager
        
        # 跟踪状态
        self.person_states: Dict[int, PersonState] = {}
        self.person_positions: Dict[int, List[Tuple[int, int]]] = defaultdict(list)
        self.max_position_history = 10
        
        # 计数数据
        self.enter_count = 0
        self.exit_count = 0
        self.current_count = 0
        
        # 记录历史
        self.count_records: List[CountRecord] = []
        self.max_records = 1000
        
        # 防抖动设置
        self.state_change_threshold = 3  # 需要连续几帧确认状态变化
        self.state_change_buffer: Dict[int, List[PersonState]] = defaultdict(list)
        
        # 日志
        self.logger = logging.getLogger(__name__)
        
        self.logger.info("人员计数器初始化完成")
    
    def update(self, track_infos: List[TrackInfo]) -> Dict[str, int]:
        """
        更新计数器
        
        Args:
            track_infos: 跟踪信息列表
            
        Returns:
            当前计数状态
        """
        current_track_ids = set()
        
        for track_info in track_infos:
            track_id = track_info.track_id
            center = track_info.center
            current_track_ids.add(track_id)
            
            # 更新位置历史
            self._update_position_history(track_id, center)
            
            # 判断当前位置状态
            current_state = self._get_current_state(center)
            
            # 处理状态变化
            self._process_state_change(track_id, current_state, track_info)
        
        # 清理不活跃的跟踪
        self._cleanup_inactive_tracks(current_track_ids)
        
        return self.get_count_status()
    
    def _update_position_history(self, track_id: int, position: Tuple[int, int]):
        """更新位置历史"""
        self.person_positions[track_id].append(position)
        if len(self.person_positions[track_id]) > self.max_position_history:
            self.person_positions[track_id].pop(0)
    
    def _get_current_state(self, position: Tuple[int, int]) -> PersonState:
        """获取当前位置状态"""
        if self.region_manager.is_point_in_region(position):
            return PersonState.INSIDE
        else:
            return PersonState.OUTSIDE
    
    def _process_state_change(self, track_id: int, current_state: PersonState, track_info: TrackInfo):
        """处理状态变化"""
        # 添加到缓冲区
        self.state_change_buffer[track_id].append(current_state)
        if len(self.state_change_buffer[track_id]) > self.state_change_threshold:
            self.state_change_buffer[track_id].pop(0)
        
        # 检查是否有稳定的状态变化
        if len(self.state_change_buffer[track_id]) == self.state_change_threshold:
            if all(state == current_state for state in self.state_change_buffer[track_id]):
                # 状态稳定，检查是否需要更新
                previous_state = self.person_states.get(track_id, PersonState.UNKNOWN)
                
                if previous_state != current_state:
                    self._handle_state_transition(track_id, previous_state, current_state, track_info)
                    self.person_states[track_id] = current_state
    
    def _handle_state_transition(self, track_id: int, previous_state: PersonState, 
                                current_state: PersonState, track_info: TrackInfo):
        """处理状态转换"""
        timestamp = time.time()
        
        # 处理进入
        if previous_state == PersonState.OUTSIDE and current_state == PersonState.INSIDE:
            self.enter_count += 1
            self.current_count += 1
            
            record = CountRecord(
                track_id=track_id,
                timestamp=timestamp,
                action="enter",
                position=track_info.center,
                confidence=track_info.confidence
            )
            self._add_record(record)
            
            self.logger.info(f"人员进入: ID={track_id}, 当前人数={self.current_count}")
        
        # 处理离开
        elif previous_state == PersonState.INSIDE and current_state == PersonState.OUTSIDE:
            self.exit_count += 1
            self.current_count = max(0, self.current_count - 1)
            
            record = CountRecord(
                track_id=track_id,
                timestamp=timestamp,
                action="exit",
                position=track_info.center,
                confidence=track_info.confidence
            )
            self._add_record(record)
            
            self.logger.info(f"人员离开: ID={track_id}, 当前人数={self.current_count}")
        
        # 处理初始状态（新出现的人员）
        elif previous_state == PersonState.UNKNOWN:
            if current_state == PersonState.INSIDE:
                # 新人员出现在区域内，按照需求不计数
                self.logger.info(f"新人员出现在区域内: ID={track_id}, 不计数")
            else:
                # 新人员出现在区域外，正常状态
                self.logger.info(f"新人员出现在区域外: ID={track_id}")
    
    def _add_record(self, record: CountRecord):
        """添加计数记录"""
        self.count_records.append(record)
        if len(self.count_records) > self.max_records:
            self.count_records.pop(0)
    
    def _cleanup_inactive_tracks(self, active_track_ids: Set[int]):
        """清理不活跃的跟踪"""
        inactive_ids = set(self.person_states.keys()) - active_track_ids
        
        for track_id in inactive_ids:
            # 如果不活跃的跟踪最后状态是在区域内，需要减少计数
            if self.person_states.get(track_id) == PersonState.INSIDE:
                # 按照需求，在区域内消失的人员不改变计数
                self.logger.info(f"人员在区域内消失: ID={track_id}, 人数不变")
            
            # 清理数据
            self.person_states.pop(track_id, None)
            self.person_positions.pop(track_id, None)
            self.state_change_buffer.pop(track_id, None)
    
    def get_count_status(self) -> Dict[str, int]:
        """获取计数状态"""
        return {
            "enter_count": self.enter_count,
            "exit_count": self.exit_count,
            "current_count": self.current_count,
            "active_tracks": len(self.person_states)
        }
    
    def get_recent_records(self, count: int = 10) -> List[CountRecord]:
        """获取最近的计数记录"""
        return self.count_records[-count:] if self.count_records else []
    
    def reset_counter(self):
        """重置计数器"""
        self.enter_count = 0
        self.exit_count = 0
        self.current_count = 0
        self.person_states.clear()
        self.person_positions.clear()
        self.state_change_buffer.clear()
        self.count_records.clear()
        
        self.logger.info("计数器已重置")
    
    def get_statistics(self) -> Dict:
        """获取统计信息"""
        if not self.count_records:
            return {
                "total_records": 0,
                "enter_rate": 0.0,
                "exit_rate": 0.0,
                "peak_count": 0,
                "average_stay_time": 0.0
            }
        
        # 计算进出频率
        time_span = self.count_records[-1].timestamp - self.count_records[0].timestamp
        enter_rate = self.enter_count / time_span if time_span > 0 else 0
        exit_rate = self.exit_count / time_span if time_span > 0 else 0
        
        return {
            "total_records": len(self.count_records),
            "enter_rate": enter_rate,
            "exit_rate": exit_rate,
            "peak_count": max(self.current_count, self.enter_count - self.exit_count),
            "time_span": time_span
        }


class CounterVisualizer:
    """计数器可视化器"""
    
    def __init__(self):
        self.font = cv2.FONT_HERSHEY_SIMPLEX
        self.font_scale = 0.8
        self.thickness = 2
        
        # 颜色定义
        self.colors = {
            "enter": (0, 255, 0),    # 绿色
            "exit": (0, 0, 255),     # 红色
            "current": (255, 255, 0), # 黄色
            "text": (255, 255, 255)   # 白色
        }
    
    def draw_counter_info(self, frame: np.ndarray, count_status: Dict[str, int],
                         position: Tuple[int, int] = (10, 30)) -> np.ndarray:
        """
        绘制计数信息
        
        Args:
            frame: 输入帧
            count_status: 计数状态
            position: 绘制位置
            
        Returns:
            绘制信息后的帧
        """
        x, y = position
        line_height = 35
        
        # 绘制背景
        bg_height = line_height * 4 + 20
        bg_width = 300
        overlay = frame.copy()
        cv2.rectangle(overlay, (x - 10, y - 25), (x + bg_width, y + bg_height), (0, 0, 0), -1)
        cv2.addWeighted(frame, 0.7, overlay, 0.3, 0, frame)
        
        # 绘制计数信息
        texts = [
            f"Enter Count: {count_status['enter_count']}",
            f"Exit Count: {count_status['exit_count']}",
            f"Current Count: {count_status['current_count']}",
            f"Active Tracks: {count_status['active_tracks']}"
        ]
        
        colors = [
            self.colors["enter"],
            self.colors["exit"],
            self.colors["current"],
            self.colors["text"]
        ]
        
        for i, (text, color) in enumerate(zip(texts, colors)):
            cv2.putText(
                frame, text, (x, y + i * line_height),
                self.font, self.font_scale, color, self.thickness
            )
        
        return frame
    
    def draw_person_states(self, frame: np.ndarray, track_infos: List[TrackInfo],
                          person_states: Dict[int, PersonState]) -> np.ndarray:
        """绘制人员状态"""
        for track_info in track_infos:
            track_id = track_info.track_id
            center = track_info.center
            state = person_states.get(track_id, PersonState.UNKNOWN)
            
            # 根据状态选择颜色
            if state == PersonState.INSIDE:
                color = (0, 255, 0)  # 绿色
                text = "IN"
            elif state == PersonState.OUTSIDE:
                color = (0, 0, 255)  # 红色
                text = "OUT"
            else:
                color = (128, 128, 128)  # 灰色
                text = "?"
            
            # 绘制状态标识
            cv2.circle(frame, center, 8, color, -1)
            cv2.putText(
                frame, text, (center[0] + 15, center[1] - 15),
                self.font, 0.5, color, 2
            )
        
        return frame
    
    def draw_count_records(self, frame: np.ndarray, recent_records: List[CountRecord],
                          position: Tuple[int, int] = (10, 200)) -> np.ndarray:
        """绘制最近的计数记录"""
        if not recent_records:
            return frame
        
        x, y = position
        
        # 绘制标题
        cv2.putText(
            frame, "Recent Records:", (x, y),
            self.font, 0.6, self.colors["text"], 2
        )
        
        # 绘制记录
        for i, record in enumerate(recent_records[-5:]):  # 最多显示5条
            action_color = self.colors["enter"] if record.action == "enter" else self.colors["exit"]
            text = f"ID:{record.track_id} {record.action.upper()}"
            
            cv2.putText(
                frame, text, (x, y + 25 + i * 20),
                self.font, 0.5, action_color, 1
            )
        
        return frame


if __name__ == "__main__":
    # 测试代码
    import time
    from region_manager import RegionManager, RegionConfig
    
    # 配置日志
    logging.basicConfig(level=logging.INFO)
    
    # 创建测试区域
    region_config = RegionConfig(
        points=[(200, 200), (600, 200), (600, 400), (200, 400)],
        margin={"top": 20, "bottom": 20, "left": 20, "right": 20},
        color=(0, 255, 0),
        thickness=2
    )
    
    region_manager = RegionManager(region_config)
    counter = PersonCounter(region_manager)
    visualizer = CounterVisualizer()
    
    # 创建测试帧
    test_frame = np.zeros((600, 800, 3), dtype=np.uint8)
    
    # 模拟跟踪数据
    test_tracks = [
        TrackInfo(1, (100, 100, 150, 150), (125, 125), 0.9, 0, []),
        TrackInfo(2, (300, 300, 350, 350), (325, 325), 0.8, 0, []),
    ]
    
    # 测试计数
    for i in range(10):
        # 模拟移动
        for track in test_tracks:
            x, y = track.center
            # 让第一个人员进入区域
            if track.track_id == 1 and i > 3:
                new_x = min(x + 10, 400)
                new_y = min(y + 10, 300)
            else:
                new_x = x + 5
                new_y = y + 2
            
            track.center = (new_x, new_y)
        
        # 更新计数器
        count_status = counter.update(test_tracks)
        
        # 可视化
        vis_frame = test_frame.copy()
        vis_frame = region_manager.draw_region(vis_frame)
        vis_frame = visualizer.draw_counter_info(vis_frame, count_status)
        vis_frame = visualizer.draw_person_states(vis_frame, test_tracks, counter.person_states)
        
        print(f"Step {i}: {count_status}")
        
        # 显示结果
        cv2.imshow("Counter Test", vis_frame)
        if cv2.waitKey(500) & 0xFF == ord('q'):
            break
    
    cv2.destroyAllWindows()
    
    # 输出统计信息
    print("统计信息:", counter.get_statistics()) 